import datetime
from unittest.mock import ANY
import zoneinfo

import numpy as np
from numpy import nan
import pandas as pd
from .conftest import assert_frame_equal, assert_index_equal

import pytest

import pytz

import pvlib
from pvlib import location
from pvlib.location import Location, lookup_altitude
from pvlib.solarposition import declination_spencer71
from pvlib.solarposition import equation_of_time_spencer71
from .conftest import requires_ephem


def test_location_required():
    Location(32.2, -111)


def test_location_all():
    Location(32.2, -111, 'US/Arizona', 700, 'Tucson')


@pytest.mark.parametrize(
    'tz,tz_expected', [
        pytest.param('UTC', 'UTC'),
        pytest.param('Etc/GMT+5', 'Etc/GMT+5'),
        pytest.param('US/Mountain', 'US/Mountain'),
        pytest.param('America/Phoenix', 'America/Phoenix'),
        pytest.param('Asia/Kathmandu', 'Asia/Kathmandu'),
        pytest.param('Asia/Yangon', 'Asia/Yangon'),
        pytest.param(datetime.timezone.utc, 'UTC'),
        pytest.param(zoneinfo.ZoneInfo('Etc/GMT-7'), 'Etc/GMT-7'),
        pytest.param(pytz.timezone('US/Arizona'), 'US/Arizona'),
        pytest.param(-6, 'Etc/GMT+6'),
        pytest.param(-11.0, 'Etc/GMT+11'),
        pytest.param(12, 'Etc/GMT-12'),
    ],
)
def test_location_tz(tz, tz_expected):
    loc = Location(32.2, -111, tz)
    assert isinstance(loc.pytz, datetime.tzinfo)  # Abstract base class.
    assert isinstance(loc.pytz, pytz.tzinfo.BaseTzInfo)
    assert type(loc.tz) is str
    assert loc.tz == tz_expected


def test_location_tz_update():
    loc = Location(32.2, -111, -11)
    assert loc.tz == 'Etc/GMT+11'
    assert loc.pytz == pytz.timezone('Etc/GMT+11')  # Deprecated attribute.

    # Updating Location's tz updates read-only time-zone attributes.
    loc.tz = 7
    assert loc.tz == 'Etc/GMT-7'
    assert loc.pytz == pytz.timezone('Etc/GMT-7')  # Deprecated attribute.


@pytest.mark.parametrize(
    'tz', [
        'invalid',
        'Etc/GMT+20',  # offset too large.
        20,  # offset too large.
    ]
)
def test_location_invalid_tz(tz):
    with pytest.raises(zoneinfo.ZoneInfoNotFoundError):
        Location(32.2, -111, tz)


@pytest.mark.parametrize(
    'tz', [
        -9.5,  # float with non-zero fractional part.
        b"bytes not str",
        [5],
    ]
)
def test_location_invalid_tz_type(tz):
    with pytest.raises(TypeError):
        Location(32.2, -111, tz)


def test_location_print_all():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    expected_str = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__str__() == expected_str


def test_location_print_pytz():
    tus = Location(32.2, -111, pytz.timezone('US/Arizona'), 700, 'Tucson')
    expected_str = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__str__() == expected_str


@pytest.fixture
def times():
    return pd.date_range(start='20160101T0600-0700',
                         end='20160101T1800-0700',
                         freq='3h')


def test_get_clearsky(mocker, times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    m = mocker.spy(pvlib.clearsky, 'ineichen')
    out = tus.get_clearsky(times)
    assert m.call_count == 1
    assert_index_equal(out.index, times)
    # check that values are 0 before sunrise and after sunset
    assert out.iloc[0, :].sum().sum() == 0
    assert out.iloc[-1:, :].sum().sum() == 0
    # check that values are > 0 during the day
    assert (out.iloc[1:-1, :] > 0).all().all()
    assert (out.columns.values == ['ghi', 'dni', 'dhi']).all()


def test_get_clearsky_ineichen_supply_linke(mocker):
    tus = Location(32.2, -111, 'US/Arizona', 700)
    times = pd.date_range(start='2014-06-24-0700', end='2014-06-25-0700',
                          freq='3h')
    mocker.spy(pvlib.clearsky, 'ineichen')
    out = tus.get_clearsky(times, linke_turbidity=3)
    # we only care that the LT is passed in this test
    pvlib.clearsky.ineichen.assert_called_once_with(ANY, ANY, 3, ANY, ANY)
    assert_index_equal(out.index, times)
    # check that values are 0 before sunrise and after sunset
    assert out.iloc[0:2, :].sum().sum() == 0
    assert out.iloc[-2:, :].sum().sum() == 0
    # check that values are > 0 during the day
    assert (out.iloc[2:-2, :] > 0).all().all()
    assert (out.columns.values == ['ghi', 'dni', 'dhi']).all()


def test_get_clearsky_haurwitz(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='haurwitz')
    expected = pd.DataFrame(data=np.array(
                            [[   0.        ],
                             [ 242.30085588],
                             [ 559.38247117],
                             [ 384.6873791 ],
                             [   0.        ]]),
                            columns=['ghi'],
                            index=times)
    assert_frame_equal(expected, clearsky)


def test_get_clearsky_simplified_solis(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis')
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  70.00146271,  638.01145669,  236.71136245],
               [ 101.69729217,  852.51950946,  577.1117803 ],
               [  86.1679965 ,  755.98048017,  385.59586091],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_apparent_elevation(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    solar_position = {'apparent_elevation': pd.Series(80, index=times),
                      'apparent_zenith': pd.Series(10, index=times)}
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                solar_position=solar_position)
    expected = pd.DataFrame(data=np.
        array([[  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919],
               [  131.3124497 ,  1001.14754036,  1108.14147919]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_dni_extra(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                dni_extra=1370)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  67.82281485,  618.15469596,  229.34422063],
               [  98.53217848,  825.98663808,  559.15039353],
               [  83.48619937,  732.45218243,  373.59500313],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky)


def test_get_clearsky_simplified_solis_pressure(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                pressure=95000)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  70.20556637,  635.53091983,  236.17716435],
               [ 102.08954904,  850.49502085,  576.28465815],
               [  86.46561686,  753.70744638,  384.90537859],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_simplified_solis_aod_pw(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    clearsky = tus.get_clearsky(times, model='simplified_solis',
                                aod700=0.25, precipitable_water=2.)
    expected = pd.DataFrame(data=np.
        array([[   0.        ,    0.        ,    0.        ],
               [  85.77821205,  374.58084365,  179.48483117],
               [ 143.52743364,  625.91745295,  490.06254157],
               [ 114.63275842,  506.52275195,  312.24711495],
               [   0.        ,    0.        ,    0.        ]]),
                            columns=['dhi', 'dni', 'ghi'],
                            index=times)
    expected = expected[['ghi', 'dni', 'dhi']]
    assert_frame_equal(expected, clearsky, check_less_precise=2)


def test_get_clearsky_valueerror(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    with pytest.raises(ValueError):
        tus.get_clearsky(times, model='invalid_model')


def test_from_tmy_3():
    from tests.iotools.test_tmy import TMY3_TESTFILE
    from pvlib.iotools import read_tmy3
    data, meta = read_tmy3(TMY3_TESTFILE, map_variables=True)
    loc = Location.from_tmy(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_from_tmy_2():
    from tests.iotools.test_tmy import TMY2_TESTFILE
    from pvlib.iotools import read_tmy2
    data, meta = read_tmy2(TMY2_TESTFILE)
    loc = Location.from_tmy(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_from_epw():
    from tests.iotools.test_epw import epw_testfile
    from pvlib.iotools import read_epw
    data, meta = read_epw(epw_testfile)
    loc = Location.from_epw(meta, data)
    assert loc.name is not None
    assert loc.altitude != 0
    assert loc.tz != 'UTC'
    assert_frame_equal(loc.weather, data)


def test_get_solarposition(expected_solpos, golden_mst):
    times = pd.date_range(datetime.datetime(2003, 10, 17, 12, 30, 30),
                          periods=1, freq='D', tz=golden_mst.tz)
    ephem_data = golden_mst.get_solarposition(times, temperature=11)
    ephem_data = np.round(ephem_data, 3)
    expected_solpos.index = times
    expected_solpos = np.round(expected_solpos, 3)
    assert_frame_equal(expected_solpos, ephem_data[expected_solpos.columns])


def test_get_airmass(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    airmass = tus.get_airmass(times)
    expected = pd.DataFrame(data=np.array(
                            [[        nan,         nan],
                             [ 3.61046506,  3.32072602],
                             [ 1.76470864,  1.62309115],
                             [ 2.45582153,  2.25874238],
                             [        nan,         nan]]),
                            columns=['airmass_relative', 'airmass_absolute'],
                            index=times)
    assert_frame_equal(expected, airmass)

    airmass = tus.get_airmass(times, model='young1994')
    expected = pd.DataFrame(data=np.array(
                            [[        nan,         nan],
                             [ 3.6075018 ,  3.31800056],
                             [ 1.7641033 ,  1.62253439],
                             [ 2.45413091,  2.25718744],
                             [        nan,         nan]]),
                            columns=['airmass_relative', 'airmass_absolute'],
                            index=times)
    assert_frame_equal(expected, airmass)


def test_get_airmass_valueerror(times):
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    with pytest.raises(ValueError):
        tus.get_airmass(times, model='invalid_model')


def test_Location___repr__():
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')

    expected = '\n'.join([
        'Location: ',
        '  name: Tucson',
        '  latitude: 32.2',
        '  longitude: -111',
        '  altitude: 700',
        '  tz: US/Arizona'
    ])
    assert tus.__repr__() == expected


@requires_ephem
def test_get_sun_rise_set_transit(golden):
    times = pd.DatetimeIndex(['2015-01-01 07:00:00', '2015-01-01 23:00:00'],
                             tz='MST')
    result = golden.get_sun_rise_set_transit(times, method='pyephem')
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])

    result = golden.get_sun_rise_set_transit(times, method='spa')
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])

    dayofyear = 1
    declination = declination_spencer71(dayofyear)
    eot = equation_of_time_spencer71(dayofyear)
    result = golden.get_sun_rise_set_transit(times, method='geometric',
                                             declination=declination,
                                             equation_of_time=eot)
    assert all(result.columns == ['sunrise', 'sunset', 'transit'])


def test_get_sun_rise_set_transit_valueerror(golden):
    times = pd.DatetimeIndex(['2015-01-01 07:00:00', '2015-01-01 23:00:00'],
                             tz='MST')
    with pytest.raises(ValueError):
        golden.get_sun_rise_set_transit(times, method='eyeball')


def test_extra_kwargs():
    with pytest.raises(TypeError, match='arbitrary_kwarg'):
        Location(32.2, -111, arbitrary_kwarg='value')


@pytest.mark.parametrize('lat,lon,expected_alt', [
    pytest.param(32.2540, -110.9742, 724, id='Tucson, USA'),
    pytest.param(-15.3875, 28.3228, 1253, id='Lusaka, Zambia'),
    pytest.param(35.6762, 139.6503, 40, id='Tokyo, Japan'),
    pytest.param(-35.2802, 149.1310, 566, id='Canberra, Australia'),
    pytest.param(4.7110, -74.0721, 2555, id='Bogota, Colombia'),
    pytest.param(31.525849, 35.449214, -415, id='Dead Sea, West Bank'),
    pytest.param(28.6139, 77.2090, 214, id='New Delhi, India'),
    pytest.param(0, 0, 0, id='Null Island,  Atlantic Ocean'),
])
def test_lookup_altitude(lat, lon, expected_alt):
    alt_found = lookup_altitude(lat, lon)
    assert alt_found == pytest.approx(expected_alt, abs=125)


def test_location_lookup_altitude(mocker):
    mocker.spy(location, 'lookup_altitude')
    tus = Location(32.2, -111, 'US/Arizona', 700, 'Tucson')
    location.lookup_altitude.assert_not_called()
    assert tus.altitude == 700
    location.lookup_altitude.reset_mock()

    tus = Location(32.2, -111, 'US/Arizona')
    location.lookup_altitude.assert_called_once_with(32.2, -111)
    assert tus.altitude == location.lookup_altitude(32.2, -111)
